# encoding: utf-8


"""

@author: tongzhenguo

@time: 2020/9/14 上午10:46

@desc:


"""
import arrow

now = arrow.now('Asia/Shanghai')


def is_trade_date(date, trade_cal_df):
    """返回是否是交易日"""
    values = trade_cal_df[trade_cal_df.cal_date == str(date)]["is_open"].values
    if len(values) == 0:
        return False
    return values[0] == "1"


def gen_k_date_list(start_date, end_date, trade_cal_df, **kwargs):
    """生成股票k线数据的月K时间线
    todo 后面可以生成其他时间级别的，比如日周
    """
    assert kwargs["frequency"] == "m"
    date_list = []
    while start_date < end_date:
        year = start_date[:4]
        month = start_date[4:6]
        if len(date_list) > 0 and date_list[-1][:6] == f"{year}{month}":
            start_date = get_n_day(start_date, 1)
            continue
        month_last_trade_date = get_month_last_trade_date(year, month, trade_cal_df)
        date_list.append(month_last_trade_date)
        if start_date == month_last_trade_date and month_last_trade_date not in date_list:
            date_list.append(month_last_trade_date)
        start_date = get_n_day(start_date, 1)
    return date_list


def last_trade_date(trade_cal_df):
    """最近交易日时间，18点之后更新当日数据"""
    if int(now_date_fmt("HH")) < 18:
        date = int(last_n_day(1, "str", format="YYYYMMDD"))
    else:
        date = int(now_date_fmt_8())
    while trade_cal_df[trade_cal_df.cal_date == str(date)]["is_open"].values[0] == "0":
        # 非交易日跳过
        # print("date %s is not open " % date)
        date = get_n_day(str(date), -1)
    return date


def get_month_last_trade_date(year, month, trade_cal_df):
    """某年某月的最后一个交易日时间"""
    for day in range(31, 0, -1):
        date = f"{year}{month}{day}" if day >= 10 else f"{year}{month}0{day}"
        df = trade_cal_df[trade_cal_df.cal_date == date]
        if df.shape[0] < 1:
            # print(f"error get_month_last_trade_date {date} not found in trade_cal_df")
            continue
        if df["is_open"].values[0] == "1":
            return date


def now_date():
    return now


def now_date_ts():
    return now.timestamp


def now_date_fmt(fmt="YYYY-MM-DD HH:mm:ss"):
    return now.format(fmt=fmt)


def now_date_fmt_10():
    return now_date_fmt("YYYY-MM-DD")


def now_date_fmt_8():
    return now_date_fmt("YYYYMMDD")


def now_year():
    return now.format(fmt="YYYY")


def last_year():
    return now.shift(years=-1).format(fmt="YYYY")


def now_quarter():
    """财务报表季度最晚发布时间"""
    mm = now.format(fmt="MM")
    dd = now.format("DD")
    if int(mm) <= 4:
        if int(mm) == 4 and int(dd) == 30:
            return 4
        else:
            return 3
    elif int(mm) <= 8:
        if int(mm) == 8 and int(dd) == 31:
            return 2
        else:
            return 1
    elif int(mm) <= 10:
        if int(mm) == 10 and int(dd) == 31:
            return 3
        else:
            return 2
    else:
        return 3


def quarterly_report_date():
    """季度报最晚日期"""
    mm = int(now.format(fmt="MM"))
    if mm in [5, 6, 7, 8]:
        # 一季报4月底公布完毕
        return now_year() + "0331"
    elif mm in [9, 10]:
        # 半年报7至8月底公布完毕
        return now_year() + "0630"
    elif mm in [11, 12]:
        # 三季报10月底公布完毕
        return now_year() + "0930"
    else:
        # 年报明年1月中旬至4月底公布完毕
        return last_year() + "1231"


def month_to_quarterly(mm):
    """月份对应季度"""
    if mm in (4, 5, 6):
        return 2
    elif mm in (7, 8, 9):
        return 3
    elif mm in (10, 11, 12):
        return 4
    else:
        return 1


def yearly_report_date():
    """年报最晚日期"""
    return last_year() + "1231"


def last_n_day(n, mode, **kwargs):
    if mode == "arrow":
        return now.shift(days=-n)
    elif mode == "timestamp":
        return now.shift(days=-n).timestamp
    elif mode == "str":
        if "format" in kwargs:
            return now.shift(days=-n).format(kwargs["format"])
        return now.shift(days=-n).format("YYYY-MM-DD HH:mm:ss")


def get_n_day(date8, n):
    date_arrow = arrow.get(date8, "YYYYMMDD").replace(tzinfo='local')
    target_date = date_arrow.shift(days=n).format("YYYYMMDD")
    return target_date


def last_n_year(n, mode, **kwargs):
    if mode == "arrow":
        return now.shift(years=-n)
    elif mode == "timestamp":
        return now.shift(years=-n).timestamp
    elif mode == "str":
        if "format" in kwargs:
            return now.shift(years=-n).format(kwargs["format"])
        return now.shift(years=-n).format("YYYY-MM-DD HH:mm:ss")


def last_n_hour(n, mode, **kwargs):
    if mode == "arrow":
        return now.shift(hours=-n)
    elif mode == "timestamp":
        return now.shift(hours=-n).timestamp
    elif mode == "str":
        if "format" in kwargs:
            return now.shift(hours=-n).format(kwargs["format"])
        return now.shift(hours=-n).format("YYYY-MM-DD HH:mm:ss")


def transform(date, src_mode, tar_mode, **kwargs):
    """
    各种时间类型转换，比如ts->时间字符串，时间字符串到时间
    :param date: arrow,时间戳(暂未实现),字符串
    :param src_mode:arrow，timestamp， str
    :param tar_mode:
    :param kwargs:
    :return: 如果类型不支持返回None,否则返回tar_mode对应的类型
    """
    date_arrow = None
    if src_mode == "arrow":
        date_arrow = date
    elif src_mode == "timestamp":
        pass
    elif src_mode == "str":
        if "src_format" in kwargs:
            # 默认utc时间，需要减8个小时；或者指定使用本地时区
            date_arrow = arrow.get(date, kwargs['src_format']).replace(tzinfo='local')
        else:
            pass
    if tar_mode == "arrow":
        return date_arrow
    elif tar_mode == "timestamp":
        return date_arrow.timestamp
    elif tar_mode == "str":
        if "tar_format" in kwargs:
            return date_arrow.format(kwargs["tar_format"])
        else:
            return date_arrow
